#include <math.h>
#include <sys/socket.h>
#include <netdb.h>
#include <netinet/in.h>
#include <unistd.h>
#include <net/if.h>
#include <sys/ioctl.h>
#include <arpa/inet.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <fcntl.h>
#include <netinet/tcp.h>
#include "socks5.h"
#include "netutils.h"
#include "net-addresses.h"
#include "log.h"

const char *dns_ntop(int af, const void *src, char *dst, int size)
{
	return inet_ntop(af, src, dst, size);
}

int dns_pton(int af, const char *src, void *dst)
{
	return inet_pton(af, src, dst);
}

int set_reuseport(int socket)
{
#ifdef SO_REUSEPORT
	int opt = 1;
	return setsockopt(socket, SOL_SOCKET, SO_REUSEPORT, &opt, sizeof(opt));
#endif
	return 0;
}

size_t get_sockaddr_len(struct sockaddr *addr)
{
	if (addr->sa_family == AF_INET) {
		return sizeof(struct sockaddr_in);
	} else if (addr->sa_family == AF_INET6) {
		return sizeof(struct sockaddr_in6);
	}
	return 0;
}

int bind_to_address(int socket_fd, const char *host)
{
	struct cork_ip ip;
	struct sockaddr_storage storage;
	memset(&storage, 0, sizeof(storage));

	if (cork_ip_init(&ip, host) != -1) {
		if (ip.version == 4) {
			struct sockaddr_in *addr = (struct sockaddr_in *)&storage;
			if(host)
				dns_pton(AF_INET, host, &addr->sin_addr);
			addr->sin_family = AF_INET;
			return bind(socket_fd, (struct sockaddr *)addr, 
				sizeof(struct sockaddr_in));
		} else if (ip.version == 6) {
			struct sockaddr_in6 *addr = (struct sockaddr_in6 *)&storage;
			if(host)
				dns_pton(AF_INET6, host, &addr->sin6_addr);
			addr->sin6_family = AF_INET6;
			return bind(socket_fd, (struct sockaddr *)addr, 
				sizeof(struct sockaddr_in6));
		}
	}
	return -1;
}

ssize_t get_sockaddr(char *host, char *port,
	struct sockaddr_storage *storage, int block)
{
	if(host == NULL || strlen(host) == 0)
		return -1;

	memset(storage, 0, sizeof(struct sockaddr_storage));
	struct cork_ip ip;
	if (cork_ip_init(&ip, host) != -1) {
		if (ip.version == 4) {
			struct sockaddr_in *addr = (struct sockaddr_in *)storage;
			addr->sin_family = AF_INET;

			dns_pton(AF_INET, host, &(addr->sin_addr));
			if (port)
				addr->sin_port = htons(atoi(port));
		} else if (ip.version == 6) {
			struct sockaddr_in6 *addr = (struct sockaddr_in6 *)storage;
			addr->sin6_family = AF_INET6;

			dns_pton(AF_INET6, host, &(addr->sin6_addr));
			if (port)
				addr->sin6_port = htons(atoi(port));
		}
		return 0;
	} else {
		struct addrinfo hints;
		struct addrinfo *result, *rp;

		memset(&hints, 0, sizeof(struct addrinfo));
		/* Return IPv4 and IPv6 choices */
		hints.ai_family   = AF_UNSPEC;
		/* We want a TCP socket */
		hints.ai_socktype = SOCK_STREAM;

		int err, i;

		for (i = 1; i < 8; i++) {
			err = getaddrinfo(host, port, &hints, &result);
			if ((!block || !err)) {
				break;
			} else {
				sleep(pow(2, i));
				sys_info("failed to resolve server name, "
					"wait %.0f seconds", pow(2, i));
			}
		}

		if (err != 0) {
			sys_info("getaddrinfo: %s", gai_strerror(err));
			return -1;
		}

		for (rp = result; rp != NULL; rp = rp->ai_next) {
			if (rp->ai_family == AF_INET) {
				memcpy(storage, rp->ai_addr, 
					sizeof(struct sockaddr_in));
				break;
			}
		}

		if (rp == NULL) {
			for (rp = result; rp != NULL; rp = rp->ai_next) {
				if (rp->ai_family == AF_INET6) {
					memcpy(storage, rp->ai_addr, 
						sizeof(struct sockaddr_in6));
					break;
				}
			}
		}

		if(result)
			freeaddrinfo(result);
		if (rp == NULL) {
			sys_info("failed to resolve remote addr");
			return -1;
		}

		return 0;
	}

	return -1;
}

int sockaddr_cmp(struct sockaddr_storage *addr1,
	struct sockaddr_storage *addr2, socklen_t len)
{
	struct sockaddr_in *p1_in   = (struct sockaddr_in *)addr1;
	struct sockaddr_in *p2_in   = (struct sockaddr_in *)addr2;
	struct sockaddr_in6 *p1_in6 = (struct sockaddr_in6 *)addr1;
	struct sockaddr_in6 *p2_in6 = (struct sockaddr_in6 *)addr2;
	if (p1_in->sin_family < p2_in->sin_family)
		return -1;
	if (p1_in->sin_family > p2_in->sin_family)
		return 1;

	/* compare ip4 */
	if (p1_in->sin_family == AF_INET) {
		/* just order it, ntohs not required */
		if (p1_in->sin_port < p2_in->sin_port)
			return -1;
		if (p1_in->sin_port > p2_in->sin_port)
			return 1;

		return memcmp(&p1_in->sin_addr, &p2_in->sin_addr, INET_SIZE);
	} else if (p1_in6->sin6_family == AF_INET6) {
		/* just order it, ntohs not required */
		if (p1_in6->sin6_port < p2_in6->sin6_port)
			return -1;
		if (p1_in6->sin6_port > p2_in6->sin6_port)
			return 1;

		return memcmp(&p1_in6->sin6_addr, &p2_in6->sin6_addr,
			INET6_SIZE);
	} else {
		/* unknown type, perform this comparison for sanity. */
		return memcmp(addr1, addr2, len);
	}
}

int sockaddr_cmp_addr(struct sockaddr_storage *addr1,
	struct sockaddr_storage *addr2, socklen_t len)
{
	struct sockaddr_in *p1_in   = (struct sockaddr_in *)addr1;
	struct sockaddr_in *p2_in   = (struct sockaddr_in *)addr2;
	struct sockaddr_in6 *p1_in6 = (struct sockaddr_in6 *)addr1;
	struct sockaddr_in6 *p2_in6 = (struct sockaddr_in6 *)addr2;
	if (p1_in->sin_family < p2_in->sin_family)
		return -1;
	if (p1_in->sin_family > p2_in->sin_family)
		return 1;
	/* compare ip4 */
	if (p1_in->sin_family == AF_INET) {
		return memcmp(&p1_in->sin_addr, &p2_in->sin_addr, INET_SIZE);
	} else if (p1_in6->sin6_family == AF_INET6) {
		return memcmp(&p1_in6->sin6_addr, &p2_in6->sin6_addr,
			INET6_SIZE);
	} else {
		/* eek unknown type, perform this comparison for sanity. */
		return memcmp(addr1, addr2, len);
	}
}

int setnonblocking(int fd)
{
	int flags;
	if (-1 == (flags = fcntl(fd, F_GETFL, 0))) {
		flags = 0;
	}
	return fcntl(fd, F_SETFL, flags | O_NONBLOCK);
}

int tcp_alive(int sock)
{
	int optval = 1;
	int optlen = sizeof(optval);
	if(setsockopt(sock, SOL_SOCKET,
			SO_KEEPALIVE, &optval, optlen) < 0) {
		sys_warn("Set tcp keepalive failed: %s\n",
			strerror(errno));
		return 0;
	}

	optval = 2;
	if(setsockopt(sock, SOL_TCP,
			TCP_KEEPCNT, &optval, optlen) < 0) {
		sys_warn("Set tcp_keepalive_probes failed: %s\n",
			strerror(errno));
		return 0;
	}

	optval = 5;
	if(setsockopt(sock, SOL_TCP,
			TCP_KEEPIDLE, &optval, optlen) < 0) {
		sys_warn("Set tcp_keepalive_time failed: %s\n",
			strerror(errno));
		return 0;
	}

	optval = 2;
	if(setsockopt(sock, SOL_TCP,
			TCP_KEEPINTVL, &optval, optlen) < 0) {
		sys_warn("Set tcp_keepalive_intvl failed: %s\n",
			strerror(errno));
		return 0;
	}

	return 1;
}

int create_and_connect_noblock(const char *addr, const char *port, int *ret)
{
	struct addrinfo hints;
	struct addrinfo *result, *rp;
	int s, conn_sock;

	memset(&hints, 0, sizeof(struct addrinfo));
	/* Return IPv4 and IPv6 choices */
	hints.ai_family   = AF_UNSPEC;
	/* We want a TCP socket */
	hints.ai_socktype = SOCK_STREAM;
	s = getaddrinfo(addr, port, &hints, &result);
	if (s != 0) {
		sys_err("getaddrinfo: %s", gai_strerror(s));
		return -1;
	}
	for (rp = result; rp != NULL; rp = rp->ai_next) {
		conn_sock = socket(rp->ai_family, 
			rp->ai_socktype, rp->ai_protocol);
		if (conn_sock == -1)
			continue;
#ifdef SO_NOSIGPIPE
		setsockopt(listen_sock, SOL_SOCKET, 
			SO_NOSIGPIPE, &opt, sizeof(opt));
#endif
		setnonblocking(conn_sock);
		*ret = connect(conn_sock, rp->ai_addr, rp->ai_addrlen);
		break;
	}

	if(result)
		freeaddrinfo(result);
	if (rp == NULL) {
		sys_err("connect failed\n");
		return -1;
	}

	return conn_sock;
}

int create_and_connect(const char *addr, const char *port)
{
	struct addrinfo hints;
	struct addrinfo *result, *rp;
	int s, conn_sock;

	memset(&hints, 0, sizeof(struct addrinfo));
	/* Return IPv4 and IPv6 choices */
	hints.ai_family   = AF_UNSPEC;
	/* We want a TCP socket */
	hints.ai_socktype = SOCK_STREAM;
	s = getaddrinfo(addr, port, &hints, &result);
	if (s != 0) {
		sys_err("getaddrinfo: %s", gai_strerror(s));
		return -1;
	}
	for (rp = result; rp != NULL; rp = rp->ai_next) {
		conn_sock = socket(rp->ai_family, 
			rp->ai_socktype, rp->ai_protocol);
		if (conn_sock == -1)
			continue;
#ifdef SO_NOSIGPIPE
		setsockopt(listen_sock, SOL_SOCKET, 
			SO_NOSIGPIPE, &opt, sizeof(opt));
#endif
		s = connect(conn_sock, rp->ai_addr, rp->ai_addrlen);
		if(s == 0) {
			/* We managed to connect successfully! */
			break;
		} else {
			sys_err("Connect failed: %s(%d)\n",
				strerror(errno), errno);
		}
		close(conn_sock);
	}
	if(result)
		freeaddrinfo(result);

	if (rp == NULL) {
		sys_err("connect failed\n");
		return -1;
	}

	return conn_sock;
}

int create_and_bind(const char *addr, const char *port)
{
	struct addrinfo hints;
	struct addrinfo *result, *rp;
	int s, listen_sock = -1;

	memset(&hints, 0, sizeof(struct addrinfo));
	/* Return IPv4 and IPv6 choices */
	hints.ai_family   = AF_UNSPEC;
	/* We want a TCP socket */
	hints.ai_socktype = SOCK_STREAM;
	/* for listen when addr is NULL */
	hints.ai_flags = AI_PASSIVE;

	s = getaddrinfo(addr, port, &hints, &result);
	if (s != 0) {
		sys_err("getaddrinfo: %s", gai_strerror(s));
		return -1;
	}

	for (rp = result; rp != NULL; rp = rp->ai_next) {
		listen_sock = socket(rp->ai_family, 
			rp->ai_socktype, rp->ai_protocol);
		if (listen_sock == -1)
			continue;

		int opt = 1;
		setsockopt(listen_sock, SOL_SOCKET, 
			SO_REUSEADDR, &opt, sizeof(opt));
#ifdef SO_NOSIGPIPE
		setsockopt(listen_sock, SOL_SOCKET, 
			SO_NOSIGPIPE, &opt, sizeof(opt));
#endif
		set_reuseport(listen_sock);

		s = bind(listen_sock, rp->ai_addr, rp->ai_addrlen);
		if (s == 0) {
			/* We managed to bind successfully! */
			break;
		} else {
			sys_err("Bind failed:%s(%d)\n", strerror(errno), errno);
		}

		close(listen_sock);
	}

	if(result)
		freeaddrinfo(result);

	if (rp == NULL) {
		sys_err("bind failed\n");
		return -1;
	}

	return listen_sock;
}

const char *astart = "10.0.0.0";
const char *aend = "10.255.255.255";
const char *bstart = "172.16.0.0";
const char *bend = "172.31.255.255";
const char *cstart = "192.168.0.0";
const char *cend = "192.168.255.255 ";
const char *local = "127.0.0.1";

static in_addr_t astart_addr = 0;
static in_addr_t aend_addr; 
static in_addr_t bstart_addr;
static in_addr_t bend_addr;
static in_addr_t cstart_addr;
static in_addr_t cend_addr;
static in_addr_t local_addr;

int ipv4_private_addr(in_addr_t addr)
{
	struct in_addr tmp;
	tmp.s_addr = addr;
	addr = ntohl(addr);

	if(astart_addr == 0) {
		astart_addr = ntohl(inet_addr(astart));
		aend_addr   = ntohl(inet_addr(aend));
		bstart_addr = ntohl(inet_addr(bstart));
		bend_addr   = ntohl(inet_addr(bend));
		cstart_addr = ntohl(inet_addr(cstart));
		cend_addr   = ntohl(inet_addr(cend));
	}

	if((addr >= astart_addr && addr <= aend_addr) ||
		(addr >= bstart_addr && addr <= bend_addr) ||
		(addr >= cstart_addr && addr <= cend_addr) || 
		(addr == 0) || (addr == local_addr)) {
		sys_info("%s is local address, closing\n",
			inet_ntoa(tmp));
		return 1;
	}
	return 0;
}

int host_private_addr(char *host, size_t hostlen)
{
	if(strncmp(host, "localhost", hostlen) == 0)
		return 1;
	return 0;
}

#ifdef CONNLOG
void _sys_info_addr(char *file, int line, char *fmt, struct dst *dst)
{
	if(dst->addrtype == SOCKS5_ADDRTYPE_IPV4) {
		struct sockaddr_in *addr = (struct sockaddr_in *)&dst->storage;
		printf("%s %d %s %s %d\n", file, line, fmt,
			inet_ntoa(addr->sin_addr), ntohs(addr->sin_port));
	} else if(dst->addrtype == SOCKS5_ADDRTYPE_IPV6) {
		return;
	} else if(dst->addrtype == SOCKS5_ADDRTYPE_DOMAIN) {
		printf("%s %d %s %s %d\n", file, line, fmt,
			dst->host.host, ntohs(dst->host.port));
	}
}
#endif
